local mathematics = require 'scripts.core.types.trueskill.mathematics'
local Gaussian = mathematics.Gaussian

local Variable = {}
Variable.__index = Variable

function Variable:new()
    local self = setmetatable(Gaussian(), Variable)
    self.messages = {}
    return self
end

function Variable:delta(other)
    local pi_delta = math.abs(self.pi - other.pi)
    if pi_delta == math.huge then
        return 0.
    end
    return math.max(math.abs(self.tau - other.tau), math.sqrt(pi_delta))
end

function Variable:set(val)
    local delta = self:delta(val)
    self.pi, self.tau = val.pi, val.tau
    return delta
end

function Variable:update_message(factor, pi, tau, message)
    message = message or Gaussian(nil, nil, pi, tau)
    local old_message = self[factor]
    self[factor] = message
    return self:set(self / old_message * message)
end

function Variable:update_value(factor, pi, tau, value)
    value = value or Gaussian(nil, nil, pi, tau)
    local old_message = self[factor]
    self[factor] = value * old_message / self
    return self:set(value)
end

function Variable:__index(k)
    if k == 'messages' then
        return rawget(self, k)
    end
    return rawget(self.messages, k) or Variable[k]
end

function Variable:__newindex(k, v)
    if k == 'messages' then
        rawset(self, k, v)
        return
    end
    rawset(self.messages, k, v)
end

function Variable:__tostring()
    local args = { type(self), tostring(Gaussian.__tostring(self)),
        #self.messages,
        #self.messages == 1 and '' or 's' }
    return ('<%s %s with %d connection%s>'):format(table.unpack(args))
end

setmetatable(Variable, {
    __index = Gaussian,
    __call = Variable.new,
})

local Factor = {}
Factor.__index = Factor

function Factor:new(variables)
    local self = setmetatable({}, Factor)
    self.vars = variables
    for _, var in ipairs(variables) do
        var[self] = Gaussian()
    end
    return self
end

function Factor:down()
    return 0
end

function Factor:up()
    return 0
end

function Factor:var()
    assert(#self.vars == 1)
    return self.vars[1]
end

function Factor:__tostring()
    local args = { type(self), #self.vars,
        #self.vars == 1 and '' or 's' }
    return ('<%s with %d connection%s>'):format(table.unpack(args))
end

setmetatable(Factor, {
    __call = Factor.new,
})

local PriorFactor = {}
PriorFactor.__index = PriorFactor

function PriorFactor:new(var, val, dynamic)
    local self = setmetatable(Factor:new({ var }), PriorFactor)
    self.val = val
    self.dynamic = dynamic or 0
    return self
end

function PriorFactor:down()
    local sigma = math.sqrt(self.val.sigma ^ 2 + self.dynamic ^ 2)
    local value = Gaussian(self.val.mu, sigma)
    return self:var():update_value(self, nil, nil, value)
end

setmetatable(PriorFactor, {
    __index = Factor,
    __call = PriorFactor.new,
})

local LikelihoodFactor = {}
LikelihoodFactor.__index = LikelihoodFactor

function LikelihoodFactor:new(mean_var, value_var, variance)
    local self = setmetatable(Factor({ mean_var, value_var }), LikelihoodFactor)
    self.mean = mean_var
    self.value = value_var
    self.variance = variance
    return self
end

function LikelihoodFactor:calc_a(var)
    return 1. / (1. + self.variance * var.pi)
end

function LikelihoodFactor:down()
    local msg = (self.mean / self.mean[self])
    local a = self:calc_a(msg)
    return self.value:update_message(self, a * msg.pi, a * msg.tau)
end

function LikelihoodFactor:up()
    local msg = (self.value / self.value[self])
    local a = self:calc_a(msg)
    return self.mean:update_message(self, a * msg.pi, a * msg.tau)
end

setmetatable(LikelihoodFactor, {
    __index = Factor,
    __call = LikelihoodFactor.new,
})

local SumFactor = {}
SumFactor.__index = SumFactor

function SumFactor:new(sum_var, term_vars, coeffs)
    local self = setmetatable(Factor({ sum_var, table.unpack(term_vars) }), SumFactor)
    self.sum = sum_var
    self.terms = term_vars
    self.coeffs = coeffs
    return self
end

function SumFactor:down()
    local vals = self.terms
    local msgs = {}
    for _, var in ipairs(vals) do
        table.insert(msgs, var[self])
    end
    return self:update(self.sum, vals, msgs, self.coeffs)
end

function SumFactor:up(index)
    index = index or 0
    local coeff = self.coeffs[index + 1]
    local coeffs = {}
    for x, c in ipairs(self.coeffs) do
        if x == index + 1 then
            table.insert(coeffs, 1. / coeff)
        else
            table.insert(coeffs, -c / coeff)
        end
    end
    local vals = { table.unpack(self.terms) }
    vals[index + 1] = self.sum
    local msgs = {}
    for _, var in ipairs(vals) do
        table.insert(msgs, var[self])
    end
    return self:update(self.terms[index + 1], vals, msgs, coeffs)
end

function SumFactor:update(var, vals, msgs, coeffs)
    local pi_inv = 0
    local mu = 0
    for i, val in ipairs(vals) do
        local msg = msgs[i]
        local coeff = coeffs[i]
        local div = val / msg
        mu = mu + coeff * div:mu()
        if pi_inv == math.huge then
            goto continue
        end
        pi_inv = pi_inv + coeff ^ 2 / div.pi
        ::continue::
    end
    local pi = 1. / pi_inv
    local tau = pi * mu
    return var:update_message(self, pi, tau)
end

setmetatable(SumFactor, {
    __index = Factor,
    __call = SumFactor.new,
})

local TruncateFactor = {}
TruncateFactor.__index = TruncateFactor

function TruncateFactor:new(var, v_func, w_func, draw_margin)
    local self = setmetatable(Factor({ var }), TruncateFactor)
    self.v_func = v_func
    self.w_func = w_func
    self.draw_margin = draw_margin
    return self
end

function TruncateFactor:up()
    local val = self:var()
    local msg = val[self]
    local div = val / msg
    local sqrt_pi = math.sqrt(div.pi)
    local args = { div.tau / sqrt_pi, self.draw_margin * sqrt_pi }
    local diff, draw_margin = table.unpack(args)
    local v = self.v_func(diff, draw_margin)
    local w = self.w_func(diff, draw_margin)
    local denom = (1. - w)
    local pi, tau = div.pi / denom, (div.tau + sqrt_pi * v) / denom
    return val:update_value(self, pi, tau)
end

setmetatable(TruncateFactor, {
    __index = Factor,
    __call = TruncateFactor.new,
})

return {
    Variable = Variable,
    PriorFactor = PriorFactor,
    LikelihoodFactor = LikelihoodFactor,
    SumFactor = SumFactor,
    TruncateFactor = TruncateFactor,
}
